import queue
import aiohttp
from aiortc import RTCPeerConnection, RTCSessionDescription, RTCConfiguration
from aiortc.mediastreams import MediaStreamTrack


class VideoTrack(MediaStreamTrack):
    # 定义轨道类型为视频
    kind = "video"

    def __init__(self, max_frames=100):
        # 调用父类的构造函数进行初始化
        super().__init__()
        # 创建一个队列，用于存储视频帧，最大容量为 max_frames
        self.frames = queue.Queue(maxsize=max_frames)

    async def recv(self):
        # 调用父类的 recv 方法接收视频帧
        return await super().recv()


# 用于建立 WebRTC 连接并接收视频帧
def rtc_frame(url, frame_queue):
    import asyncio
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    # 创建一个 RTCPeerConnection 对象，用于建立 WebRTC 连接、不使用 ICE 服务器
    pc = RTCPeerConnection(RTCConfiguration(iceServers=[]))
    # 创建自定义的视频轨道对象
    video_track = VideoTrack()
    # 将视频轨道添加到 RTCPeerConnection 中
    pc.addTrack(video_track)

    # 定义一个异步回调函数、当接收到轨道时触发
    @pc.on("track")
    async def on_track(track):
        # 检查接收到的轨道是否为视频轨道
        if track.kind == "video":
            print("接收到视频轨道、开始接收视频帧")
            while True:
                # 从轨道接收视频帧
                frame = await track.recv()
                # 将视频帧转换为 BGR24 格式的 NumPy 数组
                frame_bgr24 = frame.to_ndarray(format='bgr24')
                try:
                    # 尝试将 BGR24 格式的帧放入队列
                    frame_queue.put_nowait(frame_bgr24)
                except queue.Full:
                    # 队列已满、移除最早的帧
                    frame_queue.get_nowait()
                    # 再将新的帧放入队列
                    frame_queue.put_nowait(frame_bgr24)

    async def main():
        # 创建本地 SDP Offer
        offer = await pc.createOffer()
        print("已创建本地 SDP Offer")
        # 设置本地 SDP Offer
        await pc.setLocalDescription(offer)

        # 创建一个异步 HTTP 会话
        async with aiohttp.ClientSession() as session:
            print("开始向服务器发送 SDP Offer")
            # 向服务器发送 SDP Offer
            async with session.post(
                    url,
                    data=offer.sdp.encode(),
                    headers={
                        "Content-Type": "application/sdp",
                        "Content-Length": str(len(offer.sdp))
                    },
                    ssl=False
            ) as response:
                print("已接收到服务器的响应、开始处理 SDP Answer")
                # 从响应中获取服务器的 SDP Answer
                answer_sdp = await response.text()
                # 设置远程 SDP Answer
                await pc.setRemoteDescription(RTCSessionDescription(sdp=answer_sdp, type='answer'))

        try:
            while True:
                # 每隔 0.1 秒休眠一次，避免 CPU 占用过高
                await asyncio.sleep(0.1)
        except KeyboardInterrupt:
            # 捕获键盘中断异常，用于手动停止程序
            pass
        finally:
            print("开始关闭 RTCPeerConnection")
            # 关闭 RTCPeerConnection
            await pc.close()
            print("已关闭 RTCPeerConnection")

    loop.run_until_complete(main())
    loop.close()
